import pandas as pd
import numpy as np
import pywt
import matplotlib.pyplot as plt


# 数据处理函数（计算排沙量）
def process_data():
    INPUT_FILE = "ori_message.xlsx"
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []

    for year in years:
        df = excel_file.parse(year)
        df.columns = df.columns.str.strip()  # 去除列名空格
        required_columns = ['年', '月', '流量(m3/s)', '含沙量(kg/m3)']
        missing_columns = [col for col in required_columns if col not in df.columns]
        if missing_columns:
            raise ValueError(f"工作表 {year} 缺少必要列：{missing_columns}")

        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)

        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3)': 'mean'
        }).reset_index()

        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3)']
        all_monthly_data.append(grouped)

    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result['日期'] = pd.to_datetime(
        result['年'].astype(str) + '-' +
        result['月'].astype(str).str.zfill(2) + '-01'
    )
    return result[['日期', '排沙量(kg/s)']]


# 主流程
try:
    df = process_data()
except ValueError as e:
    print(f"数据处理错误: {e}")
    exit(1)

time_series = df['排沙量(kg/s)'].values  # 排沙量时间序列
dt = 1  # 时间间隔（月）
scales = np.arange(1, 61)  # 尺度范围（1-60，适配图16横轴）
wavelet = 'morl'  # Morlet小波

# 执行连续小波变换
[coefficients, _] = pywt.cwt(time_series, scales, wavelet, dt)

# 计算小波系数方差（每个尺度的方差）
wavelet_variance = np.var(coefficients, axis=1)

# 绘制方差图
plt.figure(figsize=(10, 6))
plt.plot(scales, wavelet_variance)
plt.xlabel('时间尺度')
plt.ylabel('小波方差')
plt.title('月输沙量小波方差')  # 与图16标题一致
plt.show()